import os
import json
from PIL import Image, ExifTags
import mysql.connector
from mysql.connector import Error
from datetime import datetime

class ImageProcessor:
    def __init__(self, db_config):
        self.db_config = db_config
        self.connection = None

    def connect_to_db(self):
        """建立数据库连接"""
        try:
            self.connection = mysql.connector.connect(**self.db_config)
            if self.connection.is_connected():
                print("成功连接到MySQL数据库")
                return True
        except Error as e:
            print(f"数据库连接错误: {e}")
            return False

    def disconnect_from_db(self):
        """关闭数据库连接"""
        if self.connection and self.connection.is_connected():
            self.connection.close()
            print("数据库连接已关闭")

    def check_image_exists(self, image_path):
        """检查数据库中是否已存在相同路径的图片"""
        cursor = None
        try:
            cursor = self.connection.cursor()
            sql = "SELECT COUNT(*) FROM news_photos WHERE image_path = %s"
            cursor.execute(sql, (image_path,))
            result = cursor.fetchone()
            return result[0] > 0
        except Error as e:
            print(f"检查图片存在性时出错: {e}")
            return False
        finally:
            if cursor:
                cursor.close()

    def get_exif_datetime(self, img):
        """从EXIF数据中获取拍摄日期和时间"""
        try:
            exif_data = img._getexif()
            if exif_data:
                # 转换EXIF标签ID为可读名称
                exif = {
                    ExifTags.TAGS[k]: v
                    for k, v in exif_data.items()
                    if k in ExifTags.TAGS
                }

                # 尝试获取原始拍摄时间(DateTimeOriginal)
                datetime_str = exif.get('DateTimeOriginal') or exif.get('DateTimeDigitized') or exif.get('DateTime')

                if datetime_str:
                    # 转换格式: "YYYY:MM:DD HH:MM:SS" → "YYYYMMDD" 和 "HH:MM"
                    try:
                        dt = datetime.strptime(datetime_str, "%Y:%m:%d %H:%M:%S")
                        return dt.strftime("%Y%m%d"), dt.strftime("%H:%M")
                    except ValueError:
                        pass
        except Exception:
            pass
        return None, None

    def get_image_info(self, image_path):
        """获取图片的元数据信息"""
        try:
            with Image.open(image_path) as img:
                width, height = img.size
                # 格式改为从文件扩展名获取，所以这里不再使用img.format

                # 获取EXIF中的拍摄日期和时间
                exif_date, exif_time = self.get_exif_datetime(img)

                # 处理DPI信息
                dpi = 72  # 默认值
                if hasattr(img, 'info') and 'dpi' in img.info:
                    try:
                        dpi_info = img.info['dpi']
                        if isinstance(dpi_info, tuple):
                            dpi = int(dpi_info[0]) if dpi_info[0] else 72
                        else:
                            dpi = int(dpi_info) if dpi_info else 72
                    except (TypeError, ValueError):
                        pass

                return {
                    "scale": f"{width}x{height}",
                    "dpi": dpi,
                    "width": width,
                    "height": height,
                    "exif_date": exif_date,
                    "exif_time": exif_time
                }
        except Exception as e:
            print(f"无法读取图片信息: {image_path}, 错误: {e}")
            return {
                "scale": "未知",
                "dpi": 72,
                "width": 0,
                "height": 0,
                "exif_date": None,
                "exif_time": None
            }

    def extract_date_from_path(self, image_path):
        """从路径中提取日期信息"""
        path_parts = image_path.split(os.sep)
        for part in path_parts:
            # 检查类似20190918的格式
            if part.isdigit() and len(part) == 8:
                try:
                    datetime.strptime(part, "%Y%m%d")
                    return part
                except ValueError:
                    continue
        return ""

    def generate_image_data(self, image_path):
        """为单个图片生成数据字典"""
        # 获取文件基本信息
        file_name = os.path.splitext(os.path.basename(image_path))[0]
        file_ext = os.path.splitext(image_path)[1][1:].upper()  # 获取文件扩展名并转为大写
        file_size = os.path.getsize(image_path) / 1024  # KB

        # 获取图片元数据
        image_info = self.get_image_info(image_path)

        # 优先使用EXIF中的日期，其次使用路径中的日期
        shooting_date = image_info["exif_date"] or self.extract_date_from_path(image_path)
        shooting_time = image_info["exif_time"] or ""

        # 构建数据字典
        image_data = {
            "image_path": image_path,
            "image_name": file_name,
            "file_number": "",  # f"0040-ZP·{shooting_date[:4]}-Y-机构代码-张号",
            "composite_title": "",
            "image_background": "",
            "image_title": "",
            "shooting_date": shooting_date,
            "shooting_period": shooting_time,  # 使用EXIF中的时间
            "shooting_location": "",
            "event_reason": "",
            "organization": "",
            "image_format": file_ext,  # 直接使用文件扩展名作为格式
            "image_scale": image_info["scale"],
            "image_size": f"{file_size:.2f} KB",
            "image_dpi": image_info["dpi"],
            "is_moved_to_archive": False,
            "journalists": "",
            "extended_data": json.dumps({
                "originalText": "",
                "persons": [],
                "image_dpi": image_info["dpi"],
                "image_width": image_info["width"],
                "image_height": image_info["height"],
                "exif_data": {
                    "shooting_date": shooting_date,
                    "shooting_time": shooting_time
                }
            })
        }

        return image_data

    def insert_image_data(self, image_data):
        """将图片数据插入数据库"""
        cursor = None
        try:
            cursor = self.connection.cursor()

            # 准备SQL语句
            sql = """
            INSERT INTO news_photos (
                image_path, image_name, file_number, composite_title,
                image_background, image_title, shooting_date, shooting_period,
                shooting_location, event_reason, organization, image_format,
                image_scale, image_size, image_dpi, is_moved_to_archive, journalists,
                extended_data
            ) VALUES (
                %(image_path)s, %(image_name)s, %(file_number)s, %(composite_title)s,
                %(image_background)s, %(image_title)s, %(shooting_date)s, %(shooting_period)s,
                %(shooting_location)s, %(event_reason)s, %(organization)s, %(image_format)s,
                %(image_scale)s, %(image_size)s, %(image_dpi)s, %(is_moved_to_archive)s, %(journalists)s,
                %(extended_data)s
            )
            """
            
            # 执行插入
            cursor.execute(sql, image_data)
            self.connection.commit()
            return True
        except Error as e:
            print(f"数据库插入错误: {e}")
            self.connection.rollback()
            return False
        finally:
            if cursor:
                cursor.close()
    
    def process_directory(self, directory):
        """处理目录下的所有图片文件"""
        supported_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff']
        processed_count = 0
        skipped_count = 0
        error_count = 0
        
        if not self.connect_to_db():
            return False
        
        try:
            for root, _, files in os.walk(directory):
                for file in files:
                    file_ext = os.path.splitext(file)[1].lower()
                    if file_ext in supported_extensions:
                        image_path = os.path.join(root, file)
                        try:
                            # 检查图片是否已存在
                            if self.check_image_exists(image_path):
                                print(f"跳过已存在文件: {image_path}")
                                skipped_count += 1
                                continue
                                
                            print(f"正在处理: {image_path}")
                            image_data = self.generate_image_data(image_path)
                            if self.insert_image_data(image_data):
                                processed_count += 1
                            else:
                                error_count += 1
                        except Exception as e:
                            print(f"处理图片 {image_path} 时出错: {e}")
                            error_count += 1
            
            print(f"\n处理完成: 成功 {processed_count} 条, 跳过 {skipped_count} 条, 失败 {error_count} 条")
            return True
        finally:
            self.disconnect_from_db()

if __name__ == "__main__":
    # 数据库配置 (根据您的实际配置修改)
    db_config = {
        'host': 'localhost',
        'user': 'root',
        'password': '123456',
        'database': 'wormhole',
        'port': 3306
    }
    
    # 图片目录路径
    image_directory = r"E:\相片档案"
    
    # 创建处理器并执行
    processor = ImageProcessor(db_config)
    processor.process_directory(image_directory)